import fire
from pycocotools.coco import COCO
from pycocotools import mask as cocomask
import numpy as np
import cv2


def prepare_masks(**kwargs):
    data_directory = "/home/shenshen/Programs/mc_data/"
    annotation_file_template = "{}/{}/annotation{}.json"

    TRAIN_IMAGES_DIRECTORY = "train/images/"
    TRAIN_ANNOTATIONS_PATH = "train/annotation.json"
    TRAIN_ANNOTATIONS_SMALL_PATH = "train/annotation-small.json"
    MASK_IMG_SAVE_PATH = "train/masks/"

    VAL_IMAGES_DIRECTORY = "val/images/"
    VAL_ANNOTATIONS_PATH = "val/annotation.json"
    VAL_ANNOTATIONS_SMALL_PATH = "val/annotation-small.json"
    VAL_MASK_IMG_SAVE_PATH = "val/masks/"

    print('generating train masks......')
    generate_masks(data_directory + TRAIN_ANNOTATIONS_PATH, data_directory + MASK_IMG_SAVE_PATH)
    print('generating val masks......')
    generate_masks(data_directory + VAL_ANNOTATIONS_PATH, data_directory + VAL_MASK_IMG_SAVE_PATH)


def generate_masks(src_annotations, target_image_path):
    coco = COCO(src_annotations)
    # category_ids = coco.loadCats(coco.getCatIds())
    # print(category_ids)

    # This generates a list of all `image_ids` available in the dataset
    image_ids = coco.getImgIds(catIds=coco.getCatIds())

    count = 0
    for image_id in image_ids:
        img = coco.loadImgs(image_id)[0]
        annotation_ids = coco.getAnnIds(imgIds=img['id'])
        annotations = coco.loadAnns(annotation_ids)

        # 一个图会有多个annotation对象
        mask = np.zeros([300, 300])  # 这里所有图像都是300*300
        for _idx, annotation in enumerate(annotations):
            rle = cocomask.frPyObjects(annotation['segmentation'], img['height'], img['width'])
            m = cocomask.decode(rle)  # (300, 300, 1)
            m = m.reshape((img['height'], img['width']))
            # 全是0和1的图片，执行OR运算求并集
            mask = np.logical_or(mask, m)

        filename = target_image_path + '{:012d}'.format(image_id) + '.jpg'
        cv2.imwrite(filename, mask * 255)

        if count % 10000 == 0:
            print("handled %s images" % (count, ))
        count += 1


if __name__ == '__main__':
    fire.Fire(prepare_masks)
